import typing
from fastapi import APIRouter, Depends
from common.custom_log import my_logger
from starlette.endpoints import WebSocketEndpoint
from starlette.websockets import WebSocket, WebSocketDisconnect
from common.state import get_state_manager, StateManager


class WsEcho(WebSocketEndpoint):
    encoding = 'json'  # json、bytes、text
    active_connections: dict = {}  # 目前的用户连接,用id做为key

    async def on_receive(self, websocket: WebSocket, data: typing.Any) -> None:
        print(data, websocket.path_params.get('id'))
        await websocket.send_json({'code': '200', 'msg': 'helloword'})

    def get_client_ip(self, websocket: WebSocket):
        # 获取客户端的 IP 地址
        client_ip = websocket.scope.get("client")
        if client_ip is None:
            # 如果 client 不存在，可能是因为代理或其他网络配置
            # 你可能需要检查 X-Forwarded-For 或其他头部来获取客户端 IP
            # 注意：这取决于你的代理设置和信任级别
            client_ip = websocket.scope.get("http_x_forwarded_for", "Unknown")
        else:
            client_ip = websocket.headers.get('x-forwarded-for')
        print(f"Client connected from IP: {client_ip}")
        return client_ip

    async def on_connect(self, websocket: WebSocket) -> None:
        try:
            await websocket.accept()
            my_logger.info(f'用户：{websocket.query_params.get("id")}接入了链接')
            my_logger.info(f'ws链接成功,token:{websocket.headers}')
            state_manager = get_state_manager()
            state_manager.set_state(f'user_{websocket.query_params.get("id")}', websocket)
            # self.active_connections[websocket.query_params.get("id")] = websocket

        except WebSocketDisconnect:
            await websocket.close()
            print("断开了连接")

    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:
        await websocket.close()
        # 移除链接
        del self.active_connections[websocket.query_params.get("id")]
